import os
from tabnanny import check
from typing import List

import matplotlib.pyplot as plt
import numpy as np


def read_tsp(path: str) -> List[List[float]]:
    """
    读取TSP问题的坐标数据文件
    """
    lines = open(path, 'r').readlines()
    assert 'NODE_COORD_SECTION\n' in lines
    index = lines.index('NODE_COORD_SECTION\n')
    data = lines[index + 1:-1]
    tmp = []
    for line in data:
        line = line.strip().split(' ')
        if line[0] == 'EOF':
            continue
        tmpline = []
        for x in line:
            if x == '':
                continue
            else:
                tmpline.append(float(x))
        if tmpline == []:
            continue
        tmp.append(tmpline)
    data = tmp
    return data


def read_tour(path: str) -> List[int]:
    """
    读取tour文件中的路径信息
    """
    lines = open(path, 'r').readlines()
    assert 'TOUR_SECTION\n' in lines
    index = lines.index('TOUR_SECTION\n')
    tour_data = lines[index + 1:]

    tour = []
    for line in tour_data:
        line = line.strip()
        if line == '-1' or line == 'EOF':
            break
        tour.append(int(line))

    tour = [x - 1 for x in tour]
    return tour


def compute_euclidean_distance(city1: List[float], city2: List[float]) -> float:
    """
    计算两个城市之间的欧几里得距离
    """
    return np.sqrt(sum([(x[0] - x[1]) ** 2 for x in zip(city1, city2)]))


def compute_path_length(tour: List[int], coordinates: np.ndarray) -> float:
    """
    计算给定路径的总长度
    """
    total_length = 0.0

    for i in range(len(tour) - 1):
        city1 = coordinates[tour[i]]
        city2 = coordinates[tour[i + 1]]
        total_length += compute_euclidean_distance(city1, city2)

    first_city = coordinates[tour[0]]
    last_city = coordinates[tour[-1]]
    total_length += compute_euclidean_distance(last_city, first_city)

    return total_length


def plot_tour(tour: List[int], coordinates: np.ndarray, title: str = "TSP Tour"):
    """
    绘制TSP路径图
    """
    tour_coords = coordinates[tour]

    closed_tour_coords = np.vstack([tour_coords, tour_coords[0]])

    plt.figure(figsize=(10, 8))
    plt.plot(closed_tour_coords[:, 0], closed_tour_coords[:, 1], 'o-', markersize=6, linewidth=1.5)

    for i, city_idx in enumerate(tour):
        plt.annotate(str(city_idx + 1), (coordinates[city_idx, 0], coordinates[city_idx, 1]),
                     textcoords="offset points", xytext=(0, 10), ha='center')

    plt.title(title)
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def main():
    tsp_file = '../data/st70.tsp' # 输入TSP数据文件
    tour_file = '../result/st70.opt.tour' # 输入Tour文件
    check_name = os.path.basename(tsp_file).replace('.tsp', '')

    tsp_data = read_tsp(tsp_file)
    coordinates = np.array(tsp_data)[:, 1:]

    optimal_tour = read_tour(tour_file)

    path_length = compute_path_length(optimal_tour, coordinates)

    print(f"TSP问题: {check_name}")
    print(f"城市数量: {len(coordinates)}")
    print(f"路径长度: {path_length:.2f}")
    print(f"路径序列: {' -> '.join(map(str, [x + 1 for x in optimal_tour]))}")

    plot_tour(optimal_tour, coordinates, f"check:{check_name}-{path_length:.2f}")


if __name__ == "__main__":
    main()
